package com.yl.test;

import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.map.MapUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.util.Collector;

import java.util.HashMap;
import java.util.List;

@Slf4j
public class CustomFunc extends RichFlatMapFunction<String, String> {

    // <key, count> key：测点id+批次id，count：收到的包个数
    private transient MapState<String, Integer> sumPacketReceive;

    @Override
    public void open(Configuration parameters) throws Exception {
        MapStateDescriptor<String, Integer> stateDescriptor = new MapStateDescriptor<>(
                "sum_receive_packet",
                Types.STRING,
                Types.INT);
        StateTtlConfig stateTtlConfig = StateTtlConfig
                .newBuilder(Time.minutes(30))
                .setUpdateType(StateTtlConfig.UpdateType.OnCreateAndWrite)
                .setStateVisibility(StateTtlConfig.StateVisibility.NeverReturnExpired)
                .build();
        stateDescriptor.enableTimeToLive(stateTtlConfig);
        sumPacketReceive = getRuntimeContext().getMapState(stateDescriptor);
    }

    @Override
    public void flatMap(String payload, Collector<String> out) throws Exception {
        log.info("payload: {}", payload);
        List<String> items = ListUtil.toList(payload.split(" "));
        int targetId = Integer.parseInt(items.get(0));
        int batchId = Integer.parseInt(items.get(1));
        String key = targetId + "_" + batchId;
        int packetSum = Integer.parseInt(items.get(2));
        int packetId = Integer.parseInt(items.get(3));
        if (sumPacketReceive.contains(key)) {
            sumPacketReceive.put(key, sumPacketReceive.get(key) + 1);
        } else {
            sumPacketReceive.put(key, 1);
        }
        for (String k : sumPacketReceive.keys()) {
            log.info("lastState: {}", k + "=" + sumPacketReceive.get(k));
        }
        if (sumPacketReceive.get(key) == packetSum) {
            out.collect(payload);
        }
    }
}
